#! utils.py
import numpy as np
import tensorflow as tf
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

class AverageMeter:

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

def xyxy_to_cxcywh(xyxy):
    return tf.concat([
        (xyxy[..., 2:] + xyxy[..., :2]) / 2, #! cx, cy
        (xyxy[..., 2:] - xyxy[..., :2]),     #!  w,  h
    ],axis=-1)

def cxcywh_to_xyxy(cxcywh):
    return tf.concat([
        cxcywh[..., :2] - cxcywh[..., 2:] / 2, #! xmin, ymin
        cxcywh[..., :2] + cxcywh[..., 2:] / 2, #! xmax, ymax
    ],axis=-1)

def cxcywh_to_gcxcywh(cxcywh, priors_cxcywh):
    """
        cxcywh:        [cx_t,cy_t,w_t,h_t]
        priors_cxcywh: [cx_d,cy_d,w_d,h_d]
        以下的10和5分别代表先验框cx,cy的方差经验值, 实际上的方差经验值为0.1和0.2
        原函数应该是:
            gcxcy = (tcxcy - pcxcy) / (variance_cxcy * pcxcy)
            gwh = log(twh / pwh) / variance_wh
        (cx_t - cx_d) / (w_d / 10) = gcx_t
        (cy_t - cy_d) / (h_d / 10) = gcy_t
        log(w_t / w_d) * 5 = gw_t
        log(h_t / h_d) * 5 = gh_t
    """
    #! 对边界框与相应先验框进行编码
    return tf.concat([
        (cxcywh[..., :2] - priors_cxcywh[..., :2]) / (priors_cxcywh[:, 2:] / 10), #! gcx, gcy
        tf.math.log(cxcywh[..., 2:] / priors_cxcywh[:, 2:]) * 5,                  #!  gw,  gh
    ],axis=-1)

def gcxcywh_to_cxcywh(gcxcywh, priors_cxcywh):
    """
        gcxcywh:       [gcx_p,gcy_p,gw_p,gh_p]
        priors_cxcywh: [ cx_d, cy_d, w_d, h_d]
        cx_p = gcx_p * (w_d / 10) + cx_d
        cy_p = gcy_p * (h_d / 10) + cy_d
        w_p = exp(gw_p / 5) * w_d
        h_p = exp(gh_p / 5) * h_d
    """
    return tf.concat([
        gcxcywh[..., :2] * priors_cxcywh[..., 2:] / 10 + priors_cxcywh[:,:2], #! cx, cy
        tf.math.exp(gcxcywh[..., 2:] / 5) * priors_cxcywh[:, 2:],             #!  w,  h
    ],axis=-1)

def find_jaccard_overlap(set_1, set_2):
    """
        寻找两个集合的交并比, set_1/set_2的格式`x1,y1,x2,y2`
        set_1 (N1,4)
        set_2 (N2,4)
        return (N1,N2)
    """
    intersection = find_intersection(set_1, set_2)

    areas_set_1 = (set_1[...,2] - set_1[...,0]) * (set_1[...,3] - set_1[...,1]) #! (N1)
    areas_set_2 = (set_2[...,2] - set_2[...,0]) * (set_2[...,3] - set_2[...,1]) #! (N2)

    union = areas_set_1[...,None] + areas_set_2[None] - intersection

    return intersection / union

def find_intersection(set_1, set_2):
    """
        计算两个集合相交的部分
        set_1/set_2的格式`x1,y1,x2,y2`
    """
    lower_bounds = np.maximum(set_1[:,None,:2], set_2[None,:,:2])  #! (N1,N2,2) (x1,y1)
    upper_bounds = np.minimum(set_1[:,None,2:], set_2[None,:,2:])  #! (N1,N2,2) (x2,y2)
    intersection_dims = np.maximum(upper_bounds - lower_bounds, 0) #! (N1,N2,2) ( w, h)
    return intersection_dims[...,0] * intersection_dims[...,1]     #! 面积

def coco_eval(gt_dataset, dt_annotations):
    """
        dt_annotations: 检测结果信息, 其中box要求的格式为: [x,y,w,h] [左上x,左上y,宽,高]
    """
    #! 初始化COCO对象
    coco_gt = COCO()
    coco_gt.dataset = gt_dataset
    coco_gt.createIndex()
    coco_dt = coco_gt.loadRes(dt_annotations)
    coco_eval = COCOeval(coco_gt, coco_dt, 'bbox')
    coco_eval.params.maxDets = [1, 10, 100]
    coco_eval.evaluate()
    coco_eval.accumulate()

    #! 打印显示
    p = coco_eval.params
    iStr = ' {:<18} {} @[ IoU={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.6f}'
    for m in range(len(p.maxDets)):
        aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == 'all']
        mind = [i for i, mDet in enumerate(p.maxDets) if mDet == p.maxDets[m]]
        s = coco_eval.eval['precision']
        s = s[:,:,:,aind,mind]
        if len(s[s>-1])==0:
            mean_s = -1
        else:
            mean_s = np.mean(s[s>-1])
        msg = iStr.format('Average Precision','(AP)',"0.50:0.95","all",p.maxDets[m],mean_s)
        tf.print(msg)